/*
 * Copyright (c) 2022 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "ola_coap_tls.h"
#include "ola_error.h"
#include "ola_log.h"
#include "ola_def.h"
#include "ola_tls_adapter.h"
#include "ola_base_utils.h"
#include "ola_pal_os.h"
#include "ola_sdk.h"
#include "ola_pal_socket.h"
#include "ola_coap.h"
#include "ola_random.h"

#define TCP_SEND_TIMEOUT 150000

static OlaTlsResource g_cloudTlsResource;

#define TLS_SELECT_WAIT_TIME        1000
#define TLS_DNS_RESOLVE_TIMEOUT_MS  (15 * 1000)
#define TLS_SELECT_TIMEOUT_MS       (15 * 1000)
#define TLS_READ_TIMEOUT_MS         TLS_SELECT_TIMEOUT_MS
#define TLS_SEND_TIMEOUT_MS         TLS_SELECT_TIMEOUT_MS
#define TLS_PORT_STR_MAX_LEN        5

static void TlsCoapError(OlaCoapContext *ctx, OlaCoapEvent event)
{
    if (ctx->eventHandler != NULL) {
        ctx->eventHandler(ctx, event);
    }
    ctx->errFlag = true;
}

int OlaTlsSend(OlaCoapSession *session, const OlaData *payload)
{
    if ((session == NULL) || (session->context == NULL) || (payload == NULL) ||
        (session->context->ssl == NULL) || (payload->data == NULL) || (payload->len == 0) ||
        session->context->errFlag) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    mbedtls_ssl_context *ssl = (mbedtls_ssl_context *)session->context->ssl;

    const unsigned char *buf = payload->data;
    unsigned int leftLen = payload->len;
    unsigned int sendLen = 0;
    unsigned long  long startTime = 0;

    if (OlaGetTime(&startTime) != OLA_OK) {
        OlaLogError("get time failed");
        return OLA_OS_ERR_GET_TIME;
    }

    int ret;
    bool isOvertime = false;
    while ((leftLen > 0) && !isOvertime) {
        ret = mbedtls_ssl_write((mbedtls_ssl_context*)ssl, buf + sendLen, leftLen);
        if ((ret < 0) || (ret > leftLen)) {
            OlaLogError("tls send failed [%d]", ret);
            TlsCoapError(session->context, OLA_COAP_EVENT_SOCKET_WRITE_ERROR);
            return OLA_MBEDTLS_ERR_SSL_SEND;
        }

        sendLen += ret;
        leftLen -= ret;
        ret = OlaIsOvertime(&isOvertime, startTime, TLS_SEND_TIMEOUT_MS);
        if (ret != OLA_OK) {
            OlaLogError("check overtime failed");
            return ret;
        }
    }

    if (leftLen > 0) {
        OlaLogError("tls send overtime");
        TlsCoapError(session->context, OLA_COAP_EVENT_SOCKET_WRITE_ERROR);
        return OLA_TLS_ERR_SEND_OVERTIME;
    }
    return OLA_OK;
}

static int TlsReadInner(OlaCoapContext *ctx, OlaBuffer *buf)
{
    unsigned char *bufCur = buf->buf;
    unsigned int recvLen = 0;
    /* 首先获取报文前2-5字节确定coap报文长度 */
    unsigned int leftLen = OLA_COAP_GET_TCP_PACKET_SIZE_MAX_BYTES;
    bool isOvertime = false;
    bool isFirst = true;

    unsigned long long startTime = 0;
    int ret = OlaGetTime(&startTime);
    if (ret != OLA_OK) {
        OlaLogError("get time failed");
        return ret;
    }

    while ((leftLen > 0) && !isOvertime) {
        ret = mbedtls_ssl_read((mbedtls_ssl_context *)ctx->ssl, bufCur + recvLen, leftLen);
        if (((ret < 0) && (ret != MBEDTLS_ERR_SSL_TIMEOUT) &&
            (ret != MBEDTLS_ERR_SSL_WANT_WRITE) && (ret != MBEDTLS_ERR_SSL_WANT_READ)) ||
            ((ret > 0) && (ret > leftLen))) {
            OlaLogError("tls read failed [%d]", ret);
            return OLA_MBEDTLS_ERR_SSL_READ;
        }
        if (ret > 0) {
            recvLen += ret;
            leftLen -= ret;
        }
        if (isFirst && (recvLen > OLA_COAP_GET_TCP_PACKET_SIZE_MIN_BYTES)) {
            int res = OlaGetTcpCoapPacketSizeByHeader(buf->buf, recvLen, &buf->len);
            if ((res != OLA_OK) || (buf->len < recvLen) || (buf->len > buf->size)) {
                OlaLogError("get coap size error %d", res);
                return OLA_COAP_ERR_GET_TCP_PACKET_SIZE;
            }
            leftLen = buf->len - recvLen;
            isFirst = false;
        }

        OlaLogDebug("tls read %u progress %u/%u", ret, recvLen, recvLen + leftLen);
        ret = OlaIsOvertime(&isOvertime, startTime, TLS_READ_TIMEOUT_MS);
        if (ret != OLA_OK) {
            OlaLogError("check overtime failed %d", ret);
            return ret;
        }
    }

    if (leftLen > 0) {
        OlaLogError("tls read overtime");
        return OLA_TLS_ERR_READ_OVERTIME;
    }
    return OLA_OK;
}

OlaCoapSession *OlaTlsRead(OlaCoapContext *ctx, OlaBuffer *buf)
{
    if ((ctx == NULL) || (buf == NULL) || (ctx->fd < 0) || ctx->errFlag ||
        (buf->buf == NULL) || (buf->size == 0) || (buf->size < OLA_COAP_GET_TCP_PACKET_SIZE_MAX_BYTES)) {
        OlaLogWarning("invalid param");
        return NULL;
    }
    int ret = TlsReadInner(ctx, buf);
    if (ret != OLA_OK) {
        TlsCoapError(ctx, OLA_COAP_EVENT_SOCKET_READ_ERROR);
        return NULL;
    }
    return ctx->session;
}

void OlaTlsDeinit(OlaCoapContext *ctx)
{
    if (ctx == NULL) {
        OlaLogWarning("invalid param");
        return;
    }

    if (ctx->fd >= 0) {
        (void)OlaClose(ctx->fd);
        ctx->fd = OLA_INVALID_FD;
    }

    if (ctx->ssl != NULL) {
        OlaReleaseTlsResource(&g_cloudTlsResource);
        ctx->ssl = NULL;
    }

    return ;
}

static int DnsResolve(const char *hostname, unsigned short port, OlaSockaddr *addr)
{
    OlaAddrInfo *answer = NULL;
    OlaAddrInfo hints;
    (void)memset_s(&hints, sizeof(OlaAddrInfo), 0, sizeof(OlaAddrInfo));
    int ret = OlaGetAddrInfo(hostname, NULL, &hints, &answer);
    if ((ret != 0) && (answer == NULL)) {
        OlaLogInfo("dns resolve failed %d", ret);
        return OLA_SOCKET_ERR_GET_ADDR_INFO;
    }
    OlaAddrInfo *point = NULL;
    ret = OLA_OK;
    for (point = answer; point != NULL; point = point->aiNext) {
        if (point->aiAddr == NULL) {
            continue;
        }
        if (memcpy_s(addr, sizeof(OlaSockaddr), point->aiAddr, sizeof(OlaSockaddr)) != EOK) {
            OlaLogError("memcpy failed");
            ret = OLA_SECUREC_ERR_MEMCOPY;
        }
        OlaSockaddrIn *addrIn = (OlaSockaddrIn *)addr;
        addrIn->sinPort = OlaHtons(port);
        break;
    }
    OlaFreeAddrInfo(answer);
    return ret;
}

static int OlaTcpConnect(const char *hostname, unsigned short port, OlaSockaddr *addr)
{
    unsigned long long startTime = 0;
    if (OlaGetTime(&startTime) != OLA_OK) {
        OlaLogError("get time failed");
        return OLA_OS_ERR_GET_TIME;
    }

    int ret;
    bool isOvertime = false;
    while (!isOvertime) {
        ret = DnsResolve(hostname, port, addr);
        if (ret == OLA_OK) {
            OlaLogInfo("dns resolve sucess");
            break;
        }

        ret = OlaIsOvertime(&isOvertime, startTime, TLS_DNS_RESOLVE_TIMEOUT_MS);
        if (ret != OLA_OK) {
            OlaLogError("check overtime failed");
            return ret;
        }
    }

    if (isOvertime) {
        OlaLogError("dns resolve overtime");
        return OLA_TLS_ERR_DNS_RESOLVE_TIMEOUT;
    }

    int fd = OlaSocket(OLA_SOCKET_TYPE_TCP);
    if (fd < 0) {
        OlaLogError("create tcp socket failed");
        return OLA_SOCKET_ERR_FD_INVALID;
    }
    OlaLogDebug("create fd %d", fd);
    ret = OlaConnect(fd, addr);
    if (ret != 0) {
        OlaSocketErrorCode err = OlaGetSocketError(fd);
        if (err != OLA_SOCKET_ERROR_EINPROGRESS) {
            OlaLogError("tcp connect failed, ret[%d] err[%d]", ret, err);
            (void)OlaClose(fd);
            fd = OLA_INVALID_FD;
            return OLA_SOCKET_ERR_CONNECT;
        }
    }

    return fd;
}

static int OlaTcpState(int fd)
{
    unsigned long long startTime = 0;
    if (OlaGetTime(&startTime) != OLA_OK) {
        OlaLogError("get time error");
        return OLA_OS_ERR_GET_TIME;
    }

    int ret = OLA_OK;
    bool isOvertime = false;
    OlaFdSet fdSet = {&fd, 1};
    while (!isOvertime) {
        ret = OlaSelect(NULL, &fdSet, NULL, TLS_SELECT_WAIT_TIME);
        if (ret < 0) {
            OlaSocketErrorCode code = OlaGetSocketError(fd);
            if (code != OLA_SOCKET_ERROR_NO_ERROR) {
                OlaLogError("select error [%d] socket error [%d]", ret, code);
                return OLA_SOCKET_ERR_SELECT;
            }
        } else if (ret > 0) {
            return OLA_OK;
        }
        OlaLogDebug("tcp connectting...");
        ret = OlaIsOvertime(&isOvertime, startTime, TLS_SELECT_TIMEOUT_MS);
        if (ret != OLA_OK) {
            OlaLogError("overtime check error");
            return ret;
        }
    }
    OlaLogDebug("tcp connect timeout!");
    return OLA_TLS_ERR_TCP_STATE_TIMEOUT;
}

static int CoapTlsConfig(OlaCoapContext *ctx, const char *url)
{
    (void)memset_s(&g_cloudTlsResource, sizeof(OlaTlsResource), 0, sizeof(OlaTlsResource));
    int ret = OlaTlsClientConfig(ctx->fd, &g_cloudTlsResource, "CloudClient", url, NULL);
    if (ret != OLA_OK) {
        OlaLogWarning("tls config error");
        return ret;
    }

    ctx->ssl = &g_cloudTlsResource.ssl;
    return OLA_OK;
}


static int CoapTlsSessionInit(OlaCoapContext *ctx, const OlaSockaddr *addr)
{
    if (ctx->session != NULL) {
        OlaLogError("context has invalid session");
        return OLA_COAP_ERR_INVALID_SESSION;
    }
    ctx->session = (OlaCoapSession *)OlaMalloc(sizeof(OlaCoapSession));
    if (ctx->session == NULL) {
        OlaLogError("malloc error");
        return OLA_OS_ERR_MALLOC;
    }
    (void)memset_s(ctx->session, sizeof(OlaCoapSession), 0, sizeof(OlaCoapSession));

    int ret = OLA_OK;
    do {
        if (memcpy_s(&ctx->session->addr, sizeof(OlaSockaddr), addr, sizeof(OlaSockaddr)) != EOK) {
            OlaLogError("memcpy error");
            ret = OLA_SECUREC_ERR_MEMCOPY;
            break;
        }
        ctx->session->context = ctx;
        ctx->session->overFlag = false;
        /* TLS coap不涉及msgid */
        ret = OlaRand((unsigned char *)&ctx->session->token, sizeof(ctx->session->token));
        if (ret != OLA_OK) {
            OlaLogError("get random error");
            break;
        }
        return OLA_OK;
    } while (0);

    OlaPutNullAfterFree((void **)&ctx->session);
    return ret;
}

int OlaTlsInit(OlaCoapContext *ctx, const char *url, unsigned short port)
{
    if ((url == NULL) || (ctx == NULL)) {
        OlaLogWarning("invalid param");
        return OLA_COM_ERR_INVALID_ARGUMENT;
    }

    OlaSockaddr addr;
    (void)memset_s(&addr, sizeof(OlaSockaddr), 0, sizeof(OlaSockaddr));
    int fd = OlaTcpConnect(url, port, &addr);
    if (fd < 0) {
        OlaLogWarning("tcp connect failed");
        return fd;
    }

    int ret = OlaTcpState(fd);
    if (ret != OLA_OK) {
        OlaLogWarning("tcp state failed");
        (void)OlaClose(fd);
        return ret;
    }

    ctx->fd = fd;
    ret = CoapTlsConfig(ctx, url);
    if (ret != OLA_OK) {
        OlaLogWarning("tls config failed");
        (void)OlaClose(fd);
        ctx->fd = OLA_INVALID_FD;
        return ret;
    }

    ret = CoapTlsSessionInit(ctx, &addr);
    if (ret != OLA_OK) {
        OlaLogWarning("tls session init failed");
        OlaReleaseTlsResource(&g_cloudTlsResource);
        (void)OlaClose(fd);
        ctx->fd = OLA_INVALID_FD;
        return ret;
    }

    return OLA_OK;
}

void OlaTlsSessionRelease(OlaCoapSession *session)
{
    OlaPutNullAfterFree((void **)&session);
}